package com.demo.userprofile.component.service.democode;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.springframework.core.io.ByteArrayResource;
import org.springframework.core.io.Resource;

import com.google.common.collect.Lists;
import com.google.common.primitives.Doubles;
import com.google.common.primitives.Floats;

import lombok.AllArgsConstructor;
import lombok.Data;
import ml.dmlc.xgboost4j.java.Booster;
import ml.dmlc.xgboost4j.java.DMatrix;
import ml.dmlc.xgboost4j.java.XGBoost;
import ml.dmlc.xgboost4j.java.XGBoostError;

/**
 * 标签实时预测示例代码，通过XGBoost模型实现实时预测服务
 *
 * @author userprofile_demo
 */
public class XGBoostPredict {

    // 存储特征映射关系
    private static Map<String, Integer> featureMap = new HashMap<>();

    private static Booster booster;

    // 读取并加载特征文件数据，其中第一列是特征ID，第二列是特征数值
    private static void loadFeature(List<Resource> featurePath) throws IOException {
        for (Resource input : featurePath) {
            BufferedReader reader = new BufferedReader(new InputStreamReader(input.getInputStream()));
            String line;
            while ((line = reader.readLine()) != null) {
                String[] items = line.split("\t");
                featureMap.put(items[1].toLowerCase(), Integer.valueOf(items[0]));
            }
            reader.close();
        }
    }

    @Data
    @AllArgsConstructor
    public static class FeatureElement {
        // 特征名称
        private String featureName;
        // 特征数值
        private float featureValue;
    }

    // 加载模型文件
    private static void loadModel(Resource modelPath) throws XGBoostError, IOException {
        booster = XGBoost.loadModel(modelPath.getInputStream());
    }


    public static void main(String[] args) throws IOException, XGBoostError {
        byte[] featureBytes = null; // 特征文件bytes
        byte[] modelBytes = null; // 模型文件bytes
        Resource featureResource = new ByteArrayResource(featureBytes);
        Resource modelResource = new ByteArrayResource(modelBytes);
        // 加载模型文件和特征文件
        loadFeature(Lists.newArrayList(featureResource));
        loadModel(modelResource);

        // 构建特征数据，特征数据来源来自调用方或者从实时特征库中查询获取
        List<FeatureElement> marrayOrNotFeature = Lists.newArrayList();
        // 示例: 特征名称，参考特征文件中的名称；特征值
        marrayOrNotFeature.add(new FeatureElement("特征1", 100.0f));
        marrayOrNotFeature.add(new FeatureElement("特征1", 20.0f));

        // 构建预测使用的特征数据
        List<Integer> featureIds = Lists.newArrayList();
        List<Float> featureValues = Lists.newArrayList();
        for (FeatureElement element : marrayOrNotFeature) {
            String featureName = element.getFeatureName().toLowerCase();
            float featureVal = element.getFeatureValue();
            Integer featureId = featureMap.get(featureName);
            if (featureId != null) {
                featureIds.add(featureId);
                featureValues.add(featureVal);
            }
        }

        // 通过模型预测数值
        int numColumn = featureIds.size();
        int[] colIndex = featureIds.stream().mapToInt(i -> i).toArray();
        long[] rowHeaders = new long[] {0, colIndex.length};
        float[] data = Floats.toArray(Doubles.asList(featureValues.stream().mapToDouble(i -> i).toArray()));
        DMatrix dmat = new DMatrix(rowHeaders, colIndex, data, DMatrix.SparseType.CSR, numColumn);
        // predicts即不同预测值概率值，根据概率值大小即可确定预测结果
        float[][] predicts = booster.predict(dmat, false, 0);
        dmat.dispose();
    }
}
